package org.example.weboj.service.impl;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.example.weboj.dto.AIChatRequest;
import org.example.weboj.dto.AIStreamResponse;
import org.example.weboj.service.AIStreamChatService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.stereotype.Service;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

@Service
public class AIStreamChatServiceImpl implements AIStreamChatService {

    private static final Logger logger = LoggerFactory.getLogger(AIStreamChatServiceImpl.class);

    @Value("${deepseek.api.key}")
    private String apiKey;

    @Value("${deepseek.api.url}")
    private String apiUrl;

    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;

    public AIStreamChatServiceImpl(RestTemplate restTemplate, ObjectMapper objectMapper) {
        this.restTemplate = restTemplate;
        this.objectMapper = objectMapper;
    }

    @Override
    public void streamChat(AIChatRequest request, WebSocketSession session) {
        // 输入验证
        if (request == null ||
            (request.getMessage() == null || request.getMessage().isEmpty()) &&
            (request.getCode() == null || request.getCode().isEmpty())) {
            sendErrorMessage(session, "请求必须包含消息或代码");
            return;
        }

        // 异步处理，避免阻塞WebSocket线程
        CompletableFuture.runAsync(() -> {
            try {
                performStreamRequest(request, session);
            } catch (Exception e) {
                logger.error("流式处理聊天时出错", e);
                sendErrorMessage(session, "处理请求时出错: " + e.getMessage());
            }
        });
    }

    private void performStreamRequest(AIChatRequest request, WebSocketSession session) {
        try {
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", "Bearer " + apiKey);
            headers.set("Accept", "text/event-stream");

            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "deepseek-chat");
            requestBody.put("stream", true);  // 启用流式响应

            // 构建提示信息
            StringBuilder prompt = new StringBuilder();
            if (request.getMessage() != null && !request.getMessage().isEmpty()) {
                prompt.append("User question: ").append(request.getMessage()).append("\n\n");
            }
            if (request.getCode() != null && !request.getCode().isEmpty()) {
                prompt.append("Code:\n").append(request.getCode()).append("\n\n");
            }
            if (request.getErrorMessage() != null && !request.getErrorMessage().isEmpty()) {
                prompt.append("Error message:\n").append(request.getErrorMessage()).append("\n\n");
            }
            if (request.getLanguage() != null && !request.getLanguage().isEmpty()) {
                prompt.append("Programming language: ").append(request.getLanguage()).append("\n\n");
            }

            requestBody.put("messages", new Object[]{
                new HashMap<String, String>() {{
                    put("role", "user");
                    put("content", prompt.toString());
                }}
            });

            requestBody.put("max_tokens", 1000);
            requestBody.put("temperature", 0.7);

            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

            logger.info("发送流式请求到DeepSeek API: {}", apiUrl);

            // 使用ResponseExtractor处理流式响应
            restTemplate.execute(apiUrl, HttpMethod.POST, req -> {
                req.getHeaders().addAll(entity.getHeaders());
                req.getBody().write(objectMapper.writeValueAsBytes(entity.getBody()));
            }, createResponseExtractor(session));

        } catch (HttpClientErrorException e) {
            logger.error("调用DeepSeek API时出错: {}", e.getMessage());
            sendErrorMessage(session, "AI服务错误: " + e.getMessage());
        } catch (Exception e) {
            logger.error("处理请求时出错", e);
            sendErrorMessage(session, "处理请求时出错: " + e.getMessage());
        }
    }

    private ResponseExtractor<Void> createResponseExtractor(WebSocketSession session) {
        return (ClientHttpResponse response) -> {
            StringBuilder cumulativeContent = new StringBuilder();
            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody()))) {
                String line;
                boolean hasAnyContent = false;

                while ((line = reader.readLine()) != null) {
                    if (line.startsWith("data: ")) {
                        String data = line.substring(6); // 删除 "data: " 前缀

                        // 处理最后一条消息
                        if ("[DONE]".equals(data)) {
                            // 发送完成消息
                            sendCompleteMessage(session);
                            break;
                        }

                        try {
                            // 解析JSON响应
                            Map<String, Object> chunk = objectMapper.readValue(data, Map.class);
                            Map<String, Object> choices = (Map<String, Object>) ((java.util.List<?>) chunk.get("choices")).get(0);
                            Map<String, Object> delta = (Map<String, Object>) choices.get("delta");

                            // 有些块可能不包含内容
                            if (delta.containsKey("content")) {
                                String content = (String) delta.get("content");
                                if (content != null && !content.isEmpty()) {
                                    cumulativeContent.append(content);
                                    hasAnyContent = true;

                                    // 发送部分响应
                                    sendPartialResponse(session, content);
                                }
                            }
                        } catch (Exception e) {
                            logger.error("解析流响应时出错: {}", e.getMessage());
                        }
                    }
                }

                // 如果没有接收到任何内容
                if (!hasAnyContent) {
                    sendErrorMessage(session, "从AI服务器接收到空响应");
                }
            } catch (IOException e) {
                logger.error("读取流响应时出错", e);
                sendErrorMessage(session, "读取AI响应时出错: " + e.getMessage());
            }
            return null;
        };
    }

    private void sendPartialResponse(WebSocketSession session, String content) {
        try {
            if (session.isOpen()) {
                AIStreamResponse response = AIStreamResponse.builder()
                        .content(content)
                        .done(false)
                        .success(true)
                        .build();
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
            }
        } catch (Exception e) {
            logger.error("发送部分响应时出错", e);
        }
    }

    private void sendCompleteMessage(WebSocketSession session) {
        try {
            if (session.isOpen()) {
                AIStreamResponse response = AIStreamResponse.builder()
                        .content("")
                        .done(true)
                        .success(true)
                        .build();
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
            }
        } catch (Exception e) {
            logger.error("发送完成消息时出错", e);
        }
    }

    private void sendErrorMessage(WebSocketSession session, String errorMessage) {
        try {
            if (session != null && session.isOpen()) {
                AIStreamResponse response = AIStreamResponse.builder()
                        .success(false)
                        .error(errorMessage)
                        .done(true)
                        .build();
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
            }
        } catch (IOException e) {
            logger.error("发送错误消息时出错", e);
        }
    }
}
